#include <stdio.h>
#include <cublas_v2.h>

static const char *_cudaGetErrorEnum(cublasStatus_t error)
{
    switch (error)
    {
    case CUBLAS_STATUS_SUCCESS:
        return "CUBLAS_STATUS_SUCCESS";

    case CUBLAS_STATUS_NOT_INITIALIZED:
        return "CUBLAS_STATUS_NOT_INITIALIZED";

    case CUBLAS_STATUS_ALLOC_FAILED:
        return "CUBLAS_STATUS_ALLOC_FAILED";

    case CUBLAS_STATUS_INVALID_VALUE:
        return "CUBLAS_STATUS_INVALID_VALUE";

    case CUBLAS_STATUS_ARCH_MISMATCH:
        return "CUBLAS_STATUS_ARCH_MISMATCH";

    case CUBLAS_STATUS_MAPPING_ERROR:
        return "CUBLAS_STATUS_MAPPING_ERROR";

    case CUBLAS_STATUS_EXECUTION_FAILED:
        return "CUBLAS_STATUS_EXECUTION_FAILED";

    case CUBLAS_STATUS_INTERNAL_ERROR:
        return "CUBLAS_STATUS_INTERNAL_ERROR";

    case CUBLAS_STATUS_NOT_SUPPORTED:
        return "CUBLAS_STATUS_NOT_SUPPORTED";

    case CUBLAS_STATUS_LICENSE_ERROR:
        return "CUBLAS_STATUS_LICENSE_ERROR";
    }
    return "<unknown>";
}

#define CHECK_CUDA_ERROR(call)                             \
    do                                                     \
    {                                                      \
        const cudaError_t errorCode = call;                \
        if (errorCode != cudaSuccess)                      \
        {                                                  \
            printf("CUDA Error:\n");                       \
            printf("    File:   %s\n", __FILE__);          \
            printf("    Line:   %d\n", __LINE__);          \
            printf("    Error code:     %d\n", errorCode); \
            printf("    Error text:     %s\n",             \
                   cudaGetErrorString(errorCode));         \
            exit(1);                                       \
        }                                                  \
    } while (0)

#define CHECK_CUBLAS_STATUS(call)                            \
    do                                                       \
    {                                                        \
        const cublasStatus_t statusCode = call;              \
        if (statusCode != CUBLAS_STATUS_SUCCESS)             \
        {                                                    \
            printf("CUDA Error:\n");                         \
            printf("    File:   %s\n", __FILE__);            \
            printf("    Line:   %d\n", __LINE__);            \
            printf("    Status code:     %d\n", statusCode); \
            printf("    Error text:     %s\n",               \
                   _cudaGetErrorEnum(statusCode));           \
            exit(1);                                         \
        }                                                    \
    } while (0)

#include <stdio.h>
#include <cub/cub.cuh>
#include <thrust/scan.h>
#include <thrust/execution_policy.h>

#ifdef USE_DP
    typedef double real;
#else
    typedef float real;
#endif


const int N = 1000000;
const int M = sizeof(int) * N;
const int NUM_REAPEATS = 20;
const int BLOCK_SIZE = 1024;
const int GRID_SIZE = (N - 1) / BLOCK_SIZE + 1;


__global__ void globalMemScan(real *d_x, real *d_y) {
    real *x = d_x + blockDim.x * blockIdx.x;
    const int n = blockDim.x * blockIdx.x + threadIdx.x;
    real y = 0.0;
    if (n < N) {
        for (int offset = 1; offset < blockDim.x; offset <<= 1) {
            if (threadIdx.x >= offset) y = x[threadIdx.x] + x[threadIdx.x - offset];
            __syncthreads();
            if (threadIdx.x >= offset) x[threadIdx.x] = y;
        }
        if (threadIdx.x == blockDim.x - 1) d_y[blockIdx.x] = x[threadIdx.x];
    } 
}

__global__ void addBaseValue(real *d_x, real *d_y) {
    const int n = blockDim.x * blockIdx.x + threadIdx.x;
    real y = blockIdx.x > 0 ? d_y[blockIdx.x - 1] : 0.0;
    if (n < N) {
        d_x[n] += y;
    } 
}

__global__ void sharedMemScan(real *d_x, real *d_y) {
    extern __shared__ real s_x[];
    const int n = blockDim.x * blockIdx.x + threadIdx.x;
    s_x[threadIdx.x] = n < N ? d_x[n] : 0.0;
    __syncthreads();
    real y = 0.0;
    if (n < N) {
        for (int offset = 1; offset < blockDim.x; offset <<= 1) {
            if (threadIdx.x >= offset) y = s_x[threadIdx.x] + s_x[threadIdx.x - offset];
            __syncthreads();
            if (threadIdx.x >= offset) s_x[threadIdx.x] = y;
        }
        d_x[n] = s_x[threadIdx.x];
        if (threadIdx.x == blockDim.x - 1) d_y[blockIdx.x] = s_x[threadIdx.x];
    } 
}


void scan(real *d_x, real *d_y, real *d_z, const int method) {
    switch (method)
    {
    case 0:
        globalMemScan<<<GRID_SIZE, BLOCK_SIZE>>>(d_x, d_y);
        globalMemScan<<<1, GRID_SIZE>>>(d_y, d_z);
        addBaseValue<<<GRID_SIZE, BLOCK_SIZE>>>(d_x, d_y);
        break;
    case 1:
        sharedMemScan<<<GRID_SIZE, BLOCK_SIZE, sizeof(real) * BLOCK_SIZE>>>(d_x, d_y);
        sharedMemScan<<<1, GRID_SIZE, sizeof(real) * GRID_SIZE>>>(d_y, d_z);
        addBaseValue<<<GRID_SIZE, BLOCK_SIZE>>>(d_x, d_y);
        break;
    
    default:
        break;
    }
}

void cub_scan(real *d_x, void *temp_storage, std::size_t &temp_storage_bytes) {
    cub::DeviceScan::InclusiveSum(temp_storage, temp_storage_bytes, d_x, d_x, N);
}

void timing(const real *h_x, real *d_x, real *d_y, real *d_z, real *h_ret, const int method) {

    float tSum = 0.0;
    float t2Sum = 0.0;
    float tSumVersion2 = 0.0;
    float t2SumVersion2 = 0.0;

    // Query CUB::DeviceScan temporary storage requirements
    void *temp_storage = nullptr;
    std::size_t temp_storage_bytes{};
    cub_scan(d_x, temp_storage, temp_storage_bytes);
    CHECK_CUDA_ERROR(cudaMalloc((void **)&temp_storage, temp_storage_bytes));

    for (int i=0; i<=NUM_REAPEATS; ++i) {
        CHECK_CUDA_ERROR(cudaMemcpy(d_x, h_x, M, cudaMemcpyHostToDevice));
        cudaEvent_t start, stop;
        CHECK_CUDA_ERROR(cudaEventCreate(&start));
        CHECK_CUDA_ERROR(cudaEventCreate(&stop));
        CHECK_CUDA_ERROR(cudaEventRecord(start));
        cudaEventQuery(start);

        if(method == 2){
            cub_scan(d_x, temp_storage, temp_storage_bytes);
        }else{
        scan(d_x, d_y, d_z, method);
        }

        CHECK_CUDA_ERROR(cudaEventRecord(stop));
        CHECK_CUDA_ERROR(cudaEventSynchronize(stop));
        float elapsedTime;
        CHECK_CUDA_ERROR(cudaEventElapsedTime(&elapsedTime, start, stop));
        if(i == 0) {
            // do nothing
        } else {
            tSumVersion2 += elapsedTime;
            t2SumVersion2 += elapsedTime * elapsedTime;
        }

        if(i == NUM_REAPEATS) {
            // do nothing
        } else {
            tSum += elapsedTime;
            t2Sum += elapsedTime * elapsedTime;
        }
        
        printf("%g\t", elapsedTime);
        CHECK_CUDA_ERROR(cudaEventDestroy(start));
        CHECK_CUDA_ERROR(cudaEventDestroy(stop));
    }
    printf("\n======version1======\n");
    printf("\n%g\t", tSum);
    float tAVG = tSum / NUM_REAPEATS;
    float tERR = sqrt(t2Sum / NUM_REAPEATS - tAVG * tAVG);
    printf("Time = %g +- %g ms.\n", tAVG, tERR);
    printf("\n======version2======\n");
    printf("\n%g\t", tSumVersion2);
    float tAVGVersion2 = tSumVersion2 / NUM_REAPEATS;
    float tERRVersion2 = sqrt(t2SumVersion2 / NUM_REAPEATS - tAVGVersion2 * tAVGVersion2);
    printf("Time = %g +- %g ms.\n", tAVGVersion2, tERRVersion2);
    CHECK_CUDA_ERROR(cudaMemcpy(h_ret, d_x, M, cudaMemcpyDeviceToHost));
}

int main() {
    real *h_x = new real[N];
    real *h_y = new real[N];    
    real *h_ret = new real[N];
    for (int i=0; i<N; i++) h_x[i] = 1.23;
    real *d_x, *d_y, *d_z;
    CHECK_CUDA_ERROR(cudaMalloc((void **)&d_x, M));
    CHECK_CUDA_ERROR(cudaMalloc((void **)&d_y, sizeof(real) * GRID_SIZE));
    CHECK_CUDA_ERROR(cudaMalloc((void **)&d_z, sizeof(real)));
    
    printf("using global mem:\n");
    timing(h_x, d_x, d_y, d_z, h_ret, 0);
    for (int i = N - 10; i < N; i++) printf("%f  ", h_ret[i]);
    printf("\n");
    printf("using shared mem:\n");
    timing(h_x, d_x, d_y, d_z, h_ret, 1);
    for (int i = N - 10; i < N; i++) printf("%f  ", h_ret[i]);
    printf("\n");
    printf("using thrust lib:\n");
    timing(h_x, d_x, d_y, d_z, h_ret, 2);
    for (int i = N - 10; i < N; i++) printf("%f  ", h_ret[i]);
    printf("\n");

    CHECK_CUDA_ERROR(cudaFree(d_x));
    CHECK_CUDA_ERROR(cudaFree(d_y));
    CHECK_CUDA_ERROR(cudaFree(d_z));
    delete[] h_x;
    delete[] h_y;
    delete[] h_ret;
}